import os
from torch.utils.data import Dataset  # 对数据进行预处理
from PIL import Image


class MyDataset(Dataset):  # 继承自 PyTorch 中的 Dataset 类
    """
    自定义 MyDatset
    """

    def __init__(self, data_path, transform=None):  # 构造函数
        """

        :param data_path: 图像文件列表的路径
        :param transform: 对图像的处理方法
        """
        super().__init__()
        self.transform = transform  # 图像与处理
        self.data_path = data_path  # 数据集目录（/VOC2012_small/）
        self.data = []  # 图片文件名列表

        img_path = os.path.join(data_path, 'img.txt')  # 拿到 img.txt 文件，里面存了所有图片的名字
        with open(img_path, 'r', encoding='utf-8') as f:
            for line in f.readlines():  # 读取每一行图像文件名
                line = line.strip()  # 去掉行末尾的空格和换行符，line 是图片的文件名：100000.jpg
                img_name = os.path.join(data_path, line)  # 组合出图像文件的完整路径

                img = Image.open(img_name)  # 根据文件路径打开图像文件，返回一个 PIL.Image.Image 对象
                if img.mode == 'RGB':  # 如果图像是 RGB 模式。
                    self.data.append(line)  # 将图像文件名加入到列表中

    def __getitem__(self, idx):  # 定义数据集类的索引函数，idx 表示数据集中的索引
        # take the data sample by its index
        img_path = os.path.join(self.data_path, self.data[idx])  # 组合出当前索引对应的图像路径

        # read image
        img = Image.open(img_path)  # 打开图像文件，返回一个 PIL.Image.Image 对象

        # apply the transform
        if self.transform:  # 判断是够定义了transform，如果定义了就执行下面的
            img = self.transform(img)  # 作用：原始的数据转换成我们需要的数据格式

        # return the image and inx
        dict_data = {  # 将数据转换成一个索引和数据对应的字典
            'index': idx,
            'img': img
        }

        return dict_data  # 返回上面的字典

    def __len__(self):  # 返回数据集的长度
        return len(self.data)
